﻿#include <algorithm>
#include <cassert>
#include "BVH.hpp"

BVHAccel::BVHAccel(std::vector<Object*> p, int maxPrimsInNode,
                   SplitMethod splitMethod)
    : maxPrimsInNode(std::min(255, maxPrimsInNode)), splitMethod(splitMethod),
      primitives(std::move(p))
{
    time_t start, stop;
    time(&start);
    if (primitives.empty())
        return;

    root = recursiveBuildSAH(primitives);

    time(&stop);
    double diff = difftime(stop, start);
    int hrs = (int)diff / 3600;
    int mins = ((int)diff / 60) - (hrs * 60);
    int secs = (int)diff - (hrs * 3600) - (mins * 60);

    printf(
        "\rBVH Generation complete: \nTime Taken: %i hrs, %i mins, %i secs\n\n",
        hrs, mins, secs);
}

BVHBuildNode* BVHAccel::recursiveBuild(std::vector<Object*> objects)
{
    BVHBuildNode* node = new BVHBuildNode();

    // Compute bounds of all primitives in BVH node
    Bounds3 bounds;
    for (int i = 0; i < objects.size(); ++i)
        bounds = Union(bounds, objects[i]->getBounds());
    if (objects.size() == 1) {
        // Create leaf _BVHBuildNode_
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr;
        node->right = nullptr;
        return node;
    }
    else if (objects.size() == 2) {
        node->left = recursiveBuild(std::vector{objects[0]});
        node->right = recursiveBuild(std::vector{objects[1]});

        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node;
    }
    else {
        Bounds3 centroidBounds;
        for (int i = 0; i < objects.size(); ++i)
            centroidBounds =
                Union(centroidBounds, objects[i]->getBounds().Centroid());
        int dim = centroidBounds.maxExtent();
        switch (dim) {
        case 0:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().x <
                       f2->getBounds().Centroid().x;
            });
            break;
        case 1:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().y <
                       f2->getBounds().Centroid().y;
            });
            break;
        case 2:
            std::sort(objects.begin(), objects.end(), [](auto f1, auto f2) {
                return f1->getBounds().Centroid().z <
                       f2->getBounds().Centroid().z;
            });
            break;
        }

        auto beginning = objects.begin();
        auto middling = objects.begin() + (objects.size() / 2);
        auto ending = objects.end();

        auto leftshapes = std::vector<Object*>(beginning, middling);
        auto rightshapes = std::vector<Object*>(middling, ending);

        assert(objects.size() == (leftshapes.size() + rightshapes.size()));

        node->left = recursiveBuild(leftshapes);
        node->right = recursiveBuild(rightshapes);

        node->bounds = Union(node->left->bounds, node->right->bounds);
    }

    return node;
}

BVHBuildNode* BVHAccel::recursiveBuildSAH(std::vector<Object*> objects)
{
    BVHBuildNode* node = new BVHBuildNode();

    // Compute bounds of all primitives in BVH node
    Bounds3 bounds;
    for (int i = 0; i < objects.size(); ++i)
        bounds = Union(bounds, objects[i]->getBounds());
    if (objects.size() == 1) {
        // Create leaf _BVHBuildNode_
        node->bounds = objects[0]->getBounds();
        node->object = objects[0];
        node->left = nullptr;
        node->right = nullptr;
        return node;
    }
    else if (objects.size() == 2) {
        node->left = recursiveBuild(std::vector{ objects[0] });
        node->right = recursiveBuild(std::vector{ objects[1] });

        node->bounds = Union(node->left->bounds, node->right->bounds);
        return node;
    }

    // 计算所有物体的质心边界
    Bounds3 centroidBounds;
    for (int i = 0; i < objects.size(); ++i)
        centroidBounds =  Union(centroidBounds, objects[i]->getBounds().Centroid());

    // 选择分割轴
    int dim = centroidBounds.maxExtent();

    // Initialize buckets for SAH 

    //•	SAH 代价公式：cost = Ct + (leftArea/totalArea) * leftCount * Ci + (rightArea/totalArea) * rightCount * Ci
    /*
        Ct: 经过该节点的遍历代价（通常设为1）
        Ci: 叶子节点交点测试代价（通常设为1）
        leftArea/rightArea: 左/右子包围盒表面积
        leftCount/rightCount: 左/右子物体数
    */

    const int nBuckets = 12; // Number of buckets for SAH
    struct BucktInfo
    {
        int count = 0;
        Bounds3 bounds;
    };
    float minCost = std::numeric_limits<float>::infinity();
    int bestSplitIndex = -1;
    int bestSplitAxis = dim;

    if (objects.size() >4) {
    
        for (int axis = 0; axis < 3; axis++){
            std::vector<BucktInfo> buckets(nBuckets);
            float mincoord = centroidBounds.pMin[axis];
            float maxcoord = centroidBounds.pMax[axis];
            float extent = maxcoord - mincoord + 1e-5f;
            // 分桶
            for (int i = 0; i < objects.size(); ++i) {
                float centroid = objects[i]->getBounds().Centroid()[axis];
                int bucketIndex = std::min(nBuckets - 1, static_cast<int>((centroid - mincoord) / extent * nBuckets));
                buckets[bucketIndex].count++;
                buckets[bucketIndex].bounds = Union(buckets[bucketIndex].bounds, objects[i]->getBounds());
            }

            // 计算每个分割点的SAH代价
            for (int i = 1; i < nBuckets; i++)
            {
                BucktInfo leftBucket, rightBucket;
                for (int j = 0; j < i; j++) {
                    leftBucket.count += buckets[j].count;
                    leftBucket.bounds = Union(leftBucket.bounds, buckets[j].bounds);
                }
                for (int j = i; j < nBuckets; j++) {
                    rightBucket.count += buckets[j].count;
                    rightBucket.bounds = Union(rightBucket.bounds, buckets[j].bounds);
                }
                if (leftBucket.count == 0 || rightBucket.count == 0)
                    continue;
                float leftArea = leftBucket.bounds.SurfaceArea();
                float rightArea = rightBucket.bounds.SurfaceArea();
                float totalArea = bounds.SurfaceArea();
                // SAH cost
                float cost = 1 + (leftArea / totalArea) * leftBucket.count + (rightArea / totalArea) * rightBucket.count;
                if (cost < minCost) {
                    minCost = cost;
                    bestSplitIndex = i;
                    bestSplitAxis = axis;
                }
            }
        }
    }

    std::vector<Object*> leftshapes, rightshapes;
    if (bestSplitIndex != -1)
    {
        // If a good split was found, use it
        int dim = bestSplitAxis;
        float mincoord = centroidBounds.pMin[dim];
        float maxcoord = centroidBounds.pMax[dim];
        float extent = maxcoord - mincoord + 1e-5f;
        for (auto& obj : objects) {
            float centroid = obj->getBounds().Centroid()[dim];
            int bucketIndex = std::min(nBuckets - 1, static_cast<int>((centroid - mincoord) / extent * nBuckets));
            if (bucketIndex < bestSplitIndex)
                leftshapes.push_back(obj);
            else
                rightshapes.push_back(obj);
        }
        // 防止极端情况，leftshapes 或者 rightshapes 为空
        if (leftshapes.empty() || rightshapes.empty()) {
            leftshapes = std::vector<Object*>(objects.begin(), objects.begin() + objects.size() / 2);
            rightshapes = std::vector<Object*>(objects.begin() + objects.size() / 2, objects.end());
        }
    }

    if(bestSplitIndex == -1 || leftshapes.empty() || rightshapes.empty()) {

        std::sort(objects.begin(), objects.end(), [bestSplitAxis](auto f1, auto f2) {
            return f1->getBounds().Centroid()[bestSplitAxis] <
                f2->getBounds().Centroid()[bestSplitAxis];
            });

        // If no good split was found, use the naive method
        leftshapes = std::vector<Object*>(objects.begin(), objects.begin() + objects.size() / 2);
        rightshapes = std::vector<Object*>(objects.begin() + objects.size() / 2, objects.end());
    }

        assert(objects.size() == (leftshapes.size() + rightshapes.size()));

        node->left = recursiveBuild(leftshapes);
        node->right = recursiveBuild(rightshapes);

        node->bounds = Union(node->left->bounds, node->right->bounds);

    return node;
}

Intersection BVHAccel::Intersect(const Ray& ray) const
{
    Intersection isect;
    if (!root)
        return isect;
    isect = BVHAccel::getIntersection(root, ray);
    return isect;
}

Intersection BVHAccel::getIntersection(BVHBuildNode* node, const Ray& ray) const
{
    // TODO Traverse the BVH to find intersection
    Intersection isect;

    // Check if the ray intersects the bounding box of the node  
    if (!node->bounds.IntersectP(ray, ray.direction_inv, { ray.direction.x > 0, ray.direction.y > 0, ray.direction.z > 0 }))
        return isect;

    // If the node is a leaf, check intersection with the object  
    if (node->left == nullptr && node->right == nullptr) {
        return node->object->getIntersection(ray);
    }

    // Otherwise, traverse the left and right children  
    Intersection leftIsect = getIntersection(node->left, ray);
    Intersection rightIsect = getIntersection(node->right, ray);

    // Return the closer intersection  
    if (leftIsect.happened && (!rightIsect.happened || leftIsect.distance < rightIsect.distance))
        return leftIsect;
    return rightIsect;

}